package org.example.algorithm.binarytree;

import org.example.algorithm.TreeNode;

import java.util.LinkedList;
import java.util.Queue;

//222. 完全二叉树的节点个数
public class CountNodesSolution {

    //二分查找+位运算
    public int countNodes(TreeNode root) {
        if (root == null) {
            return 0;
        }
        int level = getDepth(root) - 1;
        int low = 1 << level;
        int high = (1 << (level + 1)) - 1;
        while (low < high) {
            int mid = (high - low + 1) / 2 + low;
            if (existsNode(root, level, mid)) {
                low = mid;
            } else {
                high = mid - 1;
            }
        }
        return low;
    }

    private boolean existsNode(TreeNode node, int level, int k) {
        TreeNode cur = node;
        int bit = 1 << (level - 1);
        while (cur != null && bit > 0) {
            if ((bit & k) > 0) {
                cur = cur.right;
            } else {
                cur = cur.left;
            }
            bit >>= 1;
        }
        return cur != null;
    }

    //利用完全二叉树性质递归
    public int countNodes4(TreeNode root) {
        if (root == null) {
            return 0;
        }
        int leftDepth = getDepth(root.left);
        int rightDepth = getDepth(root.right);
        if (leftDepth == rightDepth) {
            return countNodes4(root.right) + (1 << leftDepth);
        } else {
            return countNodes4(root.left) + (1 << rightDepth);
        }
    }

    public int getDepth(TreeNode node) {
        int count = 0;
        TreeNode cur = node;
        while (cur != null) {
            count++;
            cur = cur.left;
        }
        return count;
    }

    //层序遍历方式
    public int countNodes3(TreeNode root) {
        if (root == null) {
            return 0;
        }
        Queue<TreeNode> queue = new LinkedList<>();
        queue.offer(root);
        int count = 1;
        while (!queue.isEmpty()) {
            TreeNode node = queue.poll();
            if (node.left != null) {
                queue.offer(node.left);
                count++;
            }
            if (node.right != null) {
                queue.offer(node.right);
                count++;
            }
        }
        return count;
    }

    //通用递归
    public int countNodes2(TreeNode root) {
        if (root == null) {
            return 0;
        }
        return countNodes2(root.left) + countNodes2(root.right) + 1;
    }
}
